package flowdroid.scan;

import db.DatabaseManager;
import flowdroid.algo.ReverseDFS;
import flowdroid.algo.ReverseDFSImpl;
import pojo.MethodMarks;
import pojo.MyTuple;
import soot.MethodOrMethodContext;
import soot.Scene;
import soot.SootMethod;
import soot.jimple.infoflow.android.manifest.ProcessManifest;
import soot.jimple.infoflow.util.SystemClassHandler;
import soot.jimple.toolkits.callgraph.ReachableMethods;
import tuple.TupleManager;
import tuple.data.DataKind;
import tuple.data.NewDataRule;
import tuple.entity.EntityRule;
import tuple.entity.MyEntity;

import java.io.*;
import java.util.*;
import java.util.stream.Collectors;

import static tuple.TupleManager.createMyTuple;
import static tuple.TupleManager.nonUserData;

public class MethodChecker {
    private final Set<SootMethod> permissionMethods;
    private final Set<SootMethod> coreMethods;
    public Map<SootMethod, MyTuple> methodsMap;

    public MethodChecker() {
        this.permissionMethods = new HashSet<>();
        this.coreMethods = new HashSet<>();

        this.methodsMap = new HashMap<>();
    }

    public void markMethods(TupleManager manager, boolean recordScanLog) throws Exception {
        System.out.println("Mark Methods");
        ReachableMethods reachableMethods = Scene.v().getReachableMethods();
        reachableMethods.update();
        int visited = 0;
        int dataCnt = 0;
        int entityCnt = 0;
        int permissionCnt = 0;

        Set<DataKind> globalData = new HashSet<>();
        Set<MyEntity> globalEntity = new HashSet<>();
        Set<String> globalPermission = new HashSet<>();
        ReverseDFS reverseDFS = new ReverseDFSImpl();

        File scanFile = new File(manager.dumpDir, "SCAN");

        try (BufferedWriter writer = recordScanLog ?
                new BufferedWriter(new FileWriter(scanFile))
                : new BufferedWriter(new OutputStreamWriter(System.out))) {
            for (Iterator<MethodOrMethodContext> iter = reachableMethods.listener(); iter.hasNext(); ) {
                try {
                    visited += 1;
                    SootMethod sm = iter.next().method();
                    boolean isValidMethod = isValidMethod(sm);

                    // 强制加载该方法的 body
//                sm.getDeclaringClass().setApplicationClass();
//                Scene.v().forceResolve(sm.getDeclaringClass().getName(), SootClass.BODIES);

                    MethodMarks methodMarks = new MethodMarks();
                    reverseDFS.checkSootClass(sm.getDeclaringClass(), methodMarks, 0);
                    reverseDFS.checkMethod(sm, methodMarks, 0);

                    Set<DataKind> data = methodMarks.dataKinds;
                    Set<MyEntity> entity = methodMarks.entitySet;
                    Set<String> stringMarks = methodMarks.stringMarks;

                    long userDataNum = data.stream().filter(x -> !nonUserData.contains(x)).count();

                    // record core methods
                    if (userDataNum >= 1 && isValidMethod) {
                        coreMethods.add(sm);

                        MyTuple tuple = createMyTuple(entity, data, stringMarks);
                        manager.addTuple(tuple);
                        methodsMap.put(sm, tuple);
                    }

                    dataCnt += !data.isEmpty() ? 1 : 0;
                    globalData.addAll(data);
                    entityCnt += !entity.isEmpty() ? 1 : 0;
                    globalEntity.addAll(entity);
                    permissionCnt += stringMarks.size() > 0 ? 1 : 0;
                    globalPermission.addAll(stringMarks);

                    writer.write(sm.getSignature() + ";");
                    if (!data.isEmpty()) {
                        writer.write(data + ";");
                    }
                    if (!entity.isEmpty()) {
                        writer.write(entity + ";");
                    }
                    writer.write(stringMarks + "\n");
                } catch (Exception e) {
                    e.printStackTrace();
                }
            }

            globalData.remove(DataKind.ANY);

            if (recordScanLog) {
                writer.write("visited:" + visited + "\n");
                writer.write("coreMethods size:" + coreMethods.size() + "\n");
            }
            writer.write("methods with data:" + dataCnt + "; data:" + globalData + "\n");
            writer.write("methods with entity:" + entityCnt + "; entity:" + globalEntity + "\n");
            writer.write("methods with condition:" + permissionCnt + "\n");
            writer.write("all Entity; size:" + globalEntity.size() + "; " + globalEntity + "\n");
            writer.write("all Data; size:" + globalData.size() + "; " + globalData + "\n");
            writer.write("all Permission; size:" + globalPermission.size() + "; " + globalPermission + "\n");
            writer.write("tuple for methods scan\n");
            String myTuple = manager.getMyTuple(manager.tuples);
            writer.write("\n");
            writer.write(myTuple);
            writer.write("\n");

            int tupleCnt = manager.tuples.size();

            try {
                DatabaseManager.insertScanResult(manager.apkName, manager.phase.toString(), visited, coreMethods.size(), tupleCnt, globalData.size(), globalEntity.size(), permissionCnt, manager.tuples.toString(), globalData.toString(), globalEntity.toString(), globalPermission.toString());
            } catch (Exception e) {
                e.printStackTrace();
                System.err.println("A DB error occurred. In file: MethodChecker.java, and caused by invoking insertScanResult(...)");
            }
        } catch (Exception e) {
            e.printStackTrace();
        }

        manager.clear();
    }

    private boolean isValidMethod(SootMethod sm) {
//        if (sm.isPhantom() || sm.getSignature().contains("dummy") || sm.getDeclaringClass().getName().contains("dummy"))
//            return false;
        if (sm.isPhantom()) return false;
        if (SystemClassHandler.v().isClassInSystemPackage(sm.getDeclaringClass().getName())) return false;
        if (sm.getDeclaringClass().isLibraryClass()) return false;
        return true;
    }


    public Set<SootMethod> getCoreMethods() {
        return coreMethods;
    }

    public void clear() {
        this.permissionMethods.clear();
        this.coreMethods.clear();
    }

    @Deprecated
    public Set<SootMethod> getPermissionMethods() {
        return permissionMethods;
    }

    @Deprecated
    public void processManifest(String apkFilePath, TupleManager manager) throws IOException {
        try (ProcessManifest processor = new ProcessManifest(apkFilePath)) {
            Set<String> permissions = processor.getPermissions();
            NewDataRule.ensureDataLoaded();
            EntityRule.ensureEntityLoaded();
            Set<DataKind> dataKinds = new HashSet<>();
            Set<MyEntity> myEntities = new HashSet<>();
            for (String str : permissions) {
                Set<String> collect = Arrays.stream(str.split("\\.")).collect(Collectors.toSet());
                dataKinds.addAll(NewDataRule.checkClassNameAsData(collect));
                myEntities.addAll(EntityRule.checkClassNameAsEntity(collect));
            }

            Set<DataKind> data = manager.getData();
            Set<MyEntity> entity = manager.getEntity();
            Set<String> permissions1 = manager.getPermissions();

            data.addAll(dataKinds);
            entity.addAll(myEntities);
            System.out.println(dataKinds);
            System.out.println(myEntities);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}